{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3123e7c-5544-40ae-a126-37a65652e3d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ====================\n",
    "# Built-in Python Libraries\n",
    "# ====================\n",
    "import argparse\n",
    "import json\n",
    "import logging\n",
    "import os\n",
    "import re\n",
    "import string\n",
    "from collections import Counter\n",
    "from typing import List, Dict, Any\n",
    "\n",
    "# ====================\n",
    "# Third-party Libraries\n",
    "# ====================\n",
    "import nltk\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "from rouge_score import rouge_scorer\n",
    "from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction\n",
    "from nltk.translate.meteor_score import meteor_score\n",
    "\n",
    "# ====================\n",
    "# Custom Modules\n",
    "# ====================\n",
    "import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "87bacf97-d5b4-4ebf-be53-f024d7dc1649",
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize_answer(s: str) -> str:\n",
    "    def remove_articles(text):\n",
    "        return re.sub(r'\\b(a|an|the)\\b', ' ', text)\n",
    "    def white_space_fix(text):\n",
    "        return ' '.join(text.split())\n",
    "    def remove_punc(text):\n",
    "        exclude = set(string.punctuation)\n",
    "        return ''.join(ch for ch in text if ch not in exclude)\n",
    "    def lower(text):\n",
    "        return text.lower()\n",
    "    return white_space_fix(remove_articles(remove_punc(lower(s))))\n",
    "    \n",
    "def calculate_f1_score_for_single_pair(\n",
    "    prediction: str,\n",
    "    ground_truth_str: str,\n",
    "):\n",
    "    \"\"\"\n",
    "    Calculate the F1 score, precision, and recall between a model's prediction and the ground truth answer.\n",
    "    \n",
    "    This function normalizes both the prediction and ground truth strings by:\n",
    "    - Converting to lowercase\n",
    "    - Removing punctuation\n",
    "    - Removing articles ('a', 'an', 'the')\n",
    "    - Normalizing whitespace\n",
    "    \n",
    "    Special handling is applied for binary answers ('yes', 'no') and 'noanswer' to avoid false positives.\n",
    "    \n",
    "    Args:\n",
    "        prediction (str): The model-generated answer (predicted string).\n",
    "        ground_truth_str (str): The correct answer (ground truth string).\n",
    "\n",
    "    Returns:\n",
    "        dict: A dictionary containing:\n",
    "            - 'f1' (float): Harmonic mean of precision and recall.\n",
    "            - 'precision' (float): Proportion of predicted tokens that are correct.\n",
    "            - 'recall' (float): Proportion of true tokens that were predicted.\n",
    "            Returns zeros if no overlap or mismatch in special cases.\n",
    "    \"\"\"\n",
    "    normalized_prediction = normalize_answer(prediction)\n",
    "    normalized_ground_truth = normalize_answer(ground_truth_str)\n",
    "\n",
    "    if normalized_prediction in ['yes', 'no', 'noanswer'] and normalized_prediction != normalized_ground_truth:\n",
    "        return {'f1': 0.0, 'precision': 0.0, 'recall': 0.0}\n",
    "    if normalized_ground_truth in ['yes', 'no', 'noanswer'] and normalized_prediction != normalized_ground_truth:\n",
    "        return {'f1': 0.0, 'precision': 0.0, 'recall': 0.0}\n",
    "\n",
    "    prediction_tokens = normalized_prediction.split()\n",
    "    ground_truth_tokens = normalized_ground_truth.split()\n",
    "    \n",
    "    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)\n",
    "    num_same = sum(common.values())\n",
    "    \n",
    "    if num_same == 0:\n",
    "        return {'f1': 0.0, 'precision': 0.0, 'recall': 0.0}\n",
    "    else:\n",
    "        precision = 1.0 * num_same / len(prediction_tokens)\n",
    "        recall = 1.0 * num_same / len(ground_truth_tokens)\n",
    "        if precision + recall == 0:\n",
    "            f1 = 0.0\n",
    "        else:\n",
    "            f1 = (2 * precision * recall) / (precision + recall)\n",
    "        \n",
    "        return {'f1': f1, 'precision': precision, 'recall': recall}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9db89eb0-80ac-468b-b951-d6ef454a55ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_scores(ref_path, gene_path):\n",
    "    \"\"\"\n",
    "    Calculate multiple evaluation metrics (ROUGE, METEOR, BLEU) between reference answers and generated predictions.\n",
    "    \n",
    "    This function reads two JSONL files: one containing ground truth answers (reference), \n",
    "    and another containing model-generated responses (predictions). It computes various text similarity \n",
    "    metrics at the question-answer pair level, aggregates them, and returns a summary DataFrame.\n",
    "\n",
    "    Args:\n",
    "        ref_path (str): Path to the reference data file in JSONL format.\n",
    "                        Each line should be a JSON object with 'question' and 'answer' fields.\n",
    "        gene_path (str): Path to the generated/predicted data file in JSONL format.\n",
    "                         Each line should contain 'question' and 'prediction' (string or list).\n",
    "\n",
    "    Returns:\n",
    "        pd.DataFrame: A single-row DataFrame containing averaged scores across all matched pairs.\n",
    "                      Columns include:\n",
    "                      - Accuracy (%)\n",
    "                      - ROUGE-1, ROUGE-2, ROUGE-L F1 scores\n",
    "                      - Average ROUGE score (%)\n",
    "                      - BLEU-1 to BLEU-4 scores\n",
    "                      - Average BLEU score (%)\n",
    "                      - Overall average score\n",
    "                      - F1 score (from custom function)\n",
    "\n",
    "    Note:\n",
    "        - Only questions present in both files are evaluated.\n",
    "        - Empty or '0' references or predictions are skipped or assigned zero scores.\n",
    "        - Uses NLTK and rouge-score libraries for metric computation.\n",
    "    \"\"\"\n",
    "    \n",
    "    # ======== Read Data ========\n",
    "    # Read reference file\n",
    "    ref_data = {}\n",
    "    with open(ref_path, 'r', encoding='utf-8') as f:\n",
    "        for line in f:\n",
    "            item = json.loads(line.strip())\n",
    "            question = item['question']\n",
    "            answer = item['answer']\n",
    "            ref_data[question] = answer\n",
    "\n",
    "    # Read generated file\n",
    "    gene_data = {}\n",
    "    with open(gene_path, 'r', encoding='utf-8') as f:\n",
    "        for line in f:\n",
    "            item = json.loads(line.strip())\n",
    "            question = item['question']\n",
    "            if isinstance(item['prediction'], list):\n",
    "                prediction = item['prediction'][-1]  # Take the last prediction\n",
    "            else:\n",
    "                prediction = item['prediction']\n",
    "            gene_data[question] = prediction\n",
    "\n",
    "    # ======== Initialize Components ========\n",
    "    scorer = rouge_scorer.RougeScorer([\"rouge1\", \"rouge2\", \"rougeL\"], use_stemmer=True)\n",
    "    smoothing = SmoothingFunction()\n",
    "\n",
    "    # ======== Result Storage ========\n",
    "    f1_scores = []\n",
    "    rouge_scores = []\n",
    "    meteor_scores = []\n",
    "    bleu1_scores = []\n",
    "    bleu2_scores = []\n",
    "    bleu3_scores = []\n",
    "    bleu4_scores = []\n",
    "\n",
    "    # ======== Calculate Metrics ========\n",
    "    matched_count = 0\n",
    "    total_count = 0\n",
    "\n",
    "    for question in ref_data:\n",
    "        total_count += 1\n",
    "        \n",
    "        if question not in gene_data:\n",
    "            # print(f\"Warning: Question not found in generated data: {question}\")\n",
    "            continue\n",
    "        \n",
    "        matched_count += 1\n",
    "        reference = str(ref_data[question]).lower()\n",
    "        generated = str(gene_data[question]).lower()\n",
    "        \n",
    "        # Skip empty reference answers\n",
    "        if not reference or reference == '0':\n",
    "            continue\n",
    "\n",
    "        f1_scores.append(calculate_f1_score_for_single_pair(generated, reference)[\"f1\"])\n",
    "        \n",
    "        # If generated answer is empty or '0', assign 0 to all metrics\n",
    "        if not generated or generated == '0':\n",
    "            rouge_scores.append({\n",
    "                'rouge1_f1': 0.0, 'rouge2_f1': 0.0, 'rougeL_f1': 0.0\n",
    "            })\n",
    "            meteor_scores.append(0.0)\n",
    "            bleu1_scores.append(0.0)\n",
    "            bleu2_scores.append(0.0)\n",
    "            bleu3_scores.append(0.0)\n",
    "            bleu4_scores.append(0.0)\n",
    "        else:\n",
    "            # ======== ROUGE ========\n",
    "            scores = scorer.score(reference, generated)\n",
    "            rouge_scores.append({\n",
    "                'rouge1_f1': scores['rouge1'].fmeasure,\n",
    "                'rouge2_f1': scores['rouge2'].fmeasure,\n",
    "                'rougeL_f1': scores['rougeL'].fmeasure\n",
    "            })\n",
    "\n",
    "            # ======== METEOR ========\n",
    "            meteor = meteor_score([reference.split()], generated.split())\n",
    "            meteor_scores.append(meteor)\n",
    "\n",
    "            # ======== BLEU 1-4 ========\n",
    "            ref_tokens = reference.split()\n",
    "            gen_tokens = generated.split()\n",
    "\n",
    "            bleu1 = sentence_bleu([ref_tokens], gen_tokens, weights=(1, 0, 0, 0), smoothing_function=smoothing.method1)\n",
    "            bleu2 = sentence_bleu([ref_tokens], gen_tokens, weights=(0.5, 0.5, 0, 0), smoothing_function=smoothing.method1)\n",
    "            bleu3 = sentence_bleu([ref_tokens], gen_tokens, weights=(0.33, 0.33, 0.33, 0), smoothing_function=smoothing.method1)\n",
    "            bleu4 = sentence_bleu([ref_tokens], gen_tokens, weights=(0.25, 0.25, 0.25, 0.25), smoothing_function=smoothing.method1)\n",
    "\n",
    "            bleu1_scores.append(bleu1)\n",
    "            bleu2_scores.append(bleu2)\n",
    "            bleu3_scores.append(bleu3)\n",
    "            bleu4_scores.append(bleu4)\n",
    "\n",
    "    print(f\"Total {total_count} questions, successfully matched {matched_count} questions\")\n",
    "    print(f\"Evaluated {len(rouge_scores)} valid question-answer pairs\")\n",
    "\n",
    "    # ======== Calculate Average Results ========\n",
    "    if rouge_scores:  # Ensure there's data before calculating averages\n",
    "        rouge_df = pd.DataFrame(rouge_scores)\n",
    "        \n",
    "        # Calculate average for each metric\n",
    "        rouge1_avg = rouge_df['rouge1_f1'].mean()\n",
    "        rouge2_avg = rouge_df['rouge2_f1'].mean()\n",
    "        rougeL_avg = rouge_df['rougeL_f1'].mean()\n",
    "        meteor_avg = sum(meteor_scores) / len(meteor_scores)\n",
    "        bleu1_avg = sum(bleu1_scores) / len(bleu1_scores)\n",
    "        bleu2_avg = sum(bleu2_scores) / len(bleu2_scores)\n",
    "        bleu3_avg = sum(bleu3_scores) / len(bleu3_scores)\n",
    "        bleu4_avg = sum(bleu4_scores) / len(bleu4_scores)\n",
    "        \n",
    "        # Calculate overall average score (using the same formula as original)\n",
    "        overall_avg = (rouge1_avg + rouge2_avg + rougeL_avg + bleu1_avg + bleu2_avg + bleu3_avg) / 6.0\n",
    "\n",
    "        metrics = utils.evaluate_predictions(\n",
    "            references_path=ref_path,\n",
    "            predictions_path=gene_path,\n",
    "            answer_field=\"answer\"\n",
    "        )\n",
    "        \n",
    "        # Create results DataFrame\n",
    "        results_df = pd.DataFrame(list(\n",
    "            [\n",
    "                round(metrics['accuracy']*100, 2),\n",
    "                rouge1_avg, rouge2_avg, rougeL_avg, round(100*sum([rouge1_avg, rouge2_avg, rougeL_avg])/3, 2),\n",
    "                bleu1_avg, bleu2_avg, bleu3_avg, round(100*sum([bleu1_avg, bleu2_avg, bleu3_avg])/3, 2),\n",
    "                sum([metrics['accuracy'], rouge1_avg, rouge2_avg, rougeL_avg, bleu1_avg, bleu2_avg, bleu3_avg])/7,\n",
    "                round(100*sum(f1_scores)/len(f1_scores), 2)\n",
    "            ]\n",
    "        )).round(4)\n",
    "        \n",
    "        return results_df\n",
    "    else:\n",
    "        print(\"No valid evaluation data found\")\n",
    "        return pd.DataFrame()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "dccc3583-83de-466e-ad1e-c4f9a1a45882",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total 885 questions, successfully matched 885 questions\n",
      "Evaluated 885 valid question-answer pairs\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>8</th>\n",
       "      <th>4</th>\n",
       "      <th>10</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>83.39</td>\n",
       "      <td>40.36</td>\n",
       "      <td>52.58</td>\n",
       "      <td>78.87</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      0      8      4      10\n",
       "0  83.39  40.36  52.58  78.87"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = get_scores(\n",
    "    \"output/GRIPRL_step_xxx.jsonl\",\n",
    "    \"test_data/formatted/WebQ_sample_10.jsonl\"\n",
    ")\n",
    "df.T[[0, 8, 4, 10]] "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (pytorch)",
   "language": "python",
   "name": "pytorch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
